#!/usr/bin/env python
"""Django's command-line utility for administrative tasks."""
import os
import sys

import torch
import torch.nn as nn


def main():
    """Run administrative tasks."""
    os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'my_vue.settings')
    try:
        from django.core.management import execute_from_command_line
    except ImportError as exc:
        raise ImportError(
            "Couldn't import Django. Are you sure it's installed and "
            "available on your PYTHONPATH environment variable? Did you "
            "forget to activate a virtual environment?"
        ) from exc
    execute_from_command_line(sys.argv)


class Mymodel(nn.Module):
    # corpus_num为不重复汉字数量, embedding_num, hidden_num, class_num为分类数量, bi=True进行双向lstm
    def __init__(self, corpus_num, embedding_num, hidden_num, class_num, bi=True):
        super().__init__()

        self.embedding = nn.Embedding(corpus_num, embedding_num)
        self.lstm = nn.LSTM(embedding_num, hidden_num, batch_first=True, bidirectional=bi)

        if bi:
            self.classifier = nn.Linear(hidden_num * 2, class_num)
        else:
            self.classifier = nn.Linear(hidden_num, class_num)

        self.cross_loss = nn.CrossEntropyLoss()

    # 做测试时没有batch_tag，先假设未传值
    def forward(self, batch_data, batch_tag=None):
        embedding = self.embedding(batch_data)
        out, _ = self.lstm(embedding)

        pre = self.classifier(out)
        self.pre = torch.argmax(pre, dim=-1).reshape(-1)  # 对预留值求最大值
        if batch_tag is not None:
            loss = self.cross_loss(pre.reshape(-1, pre.shape[-1]), batch_tag.reshape(-1))
            return loss


if __name__ == '__main__':
    main()
